# 导入必要的库
import requests
import json
import sqlite3
from bs4 import BeautifulSoup
from datetime import datetime
import time
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from webdriver_manager.chrome import ChromeDriverManager
from selenium.webdriver.chrome.options import Options

# 定义沪深交易所官网URL
shanghai_url = 'https://www.sse.com.cn/'
shenzhen_url = 'https://www.szse.cn/'

# 定义数据库文件路径
DB_FILE = 'stock_data.db'

# 创建数据库连接
conn = sqlite3.connect(DB_FILE)

# 创建游标对象
cursor = conn.cursor()

# 创建数据表，添加 exchange 字段来区分交易所
cursor.execute('''
CREATE TABLE IF NOT EXISTS stock_data (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    date TEXT,
    exchange TEXT,
    total_market_value REAL,
    circulating_market_value REAL,
    listed_companies INTEGER,
    average_pe_ratio REAL
)
''')

# 提交事务
conn.commit()

# 定义函数，用于获取网页内容
def get_page_content(url):
    if 'sse.com.cn' in url or 'szse.cn' in url:
        from webdriver_manager.microsoft import EdgeChromiumDriverManager
        from selenium.webdriver.edge.service import Service
        from selenium.webdriver.edge.options import Options
        service = Service(EdgeChromiumDriverManager().install())
        edge_options = Options()
        # 修改为指定Edge浏览器的二进制文件路径
        edge_options.binary_location = 'C:/Program Files (x86)/Microsoft/Edge/Application/msedge.exe'
        driver = webdriver.Edge(service=service, options=edge_options)
        try:
            driver.get(url)
            time.sleep(5)  # 等待页面加载
            content = driver.page_source
            return content
        except Exception as e:
            print(f'请求出错: {e}')
            return None
        finally:
            driver.quit()
    else:
        try:
            response = requests.get(url)
            response.raise_for_status()
            return response.text
        except requests.RequestException as e:
            print(f'请求出错: {e}')
            return None

# 定义函数，用于解析网页内容，获取数据
def parse_page_content(content, exchange):
    soup = BeautifulSoup(content, 'html.parser')
    if exchange == '上海交易所':
        # 上海证券交易所的解析逻辑
        total_market_value_elem = soup.select_one('.tab_content.js_marketData .sse_market_data_con li:has(span:contains("总市值/亿元")) .tab_data')
        if total_market_value_elem:
            total_market_value = total_market_value_elem.text
            print(f'{exchange}提取到的总市值: {total_market_value}')
        else:
            total_market_value = None
            print(f'{exchange}未提取到总市值')

        circulating_market_value_elem = soup.select_one('.tab_content.js_marketData .sse_market_data_con li:has(span:contains("流通市值/亿元")) .tab_data')
        if circulating_market_value_elem:
            circulating_market_value = circulating_market_value_elem.text
            print(f'{exchange}提取到的流通市值: {circulating_market_value}')
        else:
            circulating_market_value = None
            print(f'{exchange}未提取到流通市值')

        listed_companies_elem = soup.select_one('.tab_content.js_marketData .sse_market_data_con li:has(span:contains("上市公司/家")) .tab_data')
        if listed_companies_elem:
            listed_companies = listed_companies_elem.text
            print(f'{exchange}提取到的上市公司数量: {listed_companies}')
        else:
            listed_companies = None
            print(f'{exchange}未提取到上市公司数量')

        average_pe_ratio_elem = soup.select_one('.tab_content.js_marketData .sse_market_data_con li:has(span:contains("平均市盈率/倍")) .tab_data')
        if average_pe_ratio_elem:
            average_pe_ratio = average_pe_ratio_elem.text
            print(f'{exchange}提取到的平均市盈率: {average_pe_ratio}')
        else:
            average_pe_ratio = None
            print(f'{exchange}未提取到平均市盈率')
    elif exchange == '深圳交易所':
        # 深圳证券交易所的解析逻辑
        total_market_value_elem = soup.select_one('a[href="http://www.szse.cn/market/stock/indicator/index.html"] .list li:has(.tit:contains("股票总市值（亿元）")) .val')
        if total_market_value_elem:
            total_market_value = total_market_value_elem.text.strip()
            print(f'{exchange}提取到的总市值: {total_market_value}')
        else:
            total_market_value = None
            print(f'{exchange}未提取到总市值')

        circulating_market_value_elem = soup.select_one('a[href="http://www.szse.cn/market/stock/indicator/index.html"] .list li:has(.tit:contains("股票流通市值（亿元）")) .val')
        if circulating_market_value_elem:
            circulating_market_value = circulating_market_value_elem.text.strip()
            print(f'{exchange}提取到的流通市值: {circulating_market_value}')
        else:
            circulating_market_value = None
            print(f'{exchange}未提取到流通市值')

        listed_companies_elem = soup.select_one('a[href="http://www.szse.cn/market/stock/indicator/index.html"] .list li:has(.tit:contains("上市公司数")) .val')
        if listed_companies_elem:
            listed_companies = listed_companies_elem.text.strip()
            print(f'{exchange}提取到的上市公司数量: {listed_companies}')
        else:
            listed_companies = None
            print(f'{exchange}未提取到上市公司数量')

        average_pe_ratio_elem = soup.select_one('a[href="http://www.szse.cn/market/stock/indicator/index.html"] .list li:has(.tit:contains("股票平均市盈率")) .val')
        if average_pe_ratio_elem:
            average_pe_ratio = average_pe_ratio_elem.text.strip()
            print(f'{exchange}提取到的平均市盈率: {average_pe_ratio}')
        else:
            average_pe_ratio = None
            print(f'{exchange}未提取到平均市盈率')
    return total_market_value, circulating_market_value, listed_companies, average_pe_ratio

# 检查数据保存到数据库的代码，确保数据能正确写入
def save_data_to_db(date, exchange, total_market_value, circulating_market_value, listed_companies, average_pe_ratio):
    try:
        cursor.execute('''
        INSERT INTO stock_data (date, exchange, total_market_value, circulating_market_value, listed_companies, average_pe_ratio)
        VALUES (?,?,?,?,?,?)
        ''', (date, exchange, total_market_value, circulating_market_value, listed_companies, average_pe_ratio))
        conn.commit()
        print('数据保存到数据库成功')
    except Exception as e:
        print(f'数据保存到数据库失败: {e}')
        conn.rollback()
    return total_market_value, circulating_market_value, listed_companies, average_pe_ratio

# 主函数
def main():
    # 获取当前日期
    today = datetime.now().strftime('%Y-%m-%d')

    # 获取沪深交易所官网网页内容
    shanghai_content = get_page_content(shanghai_url)
    shenzhen_content = get_page_content(shenzhen_url)

    if shanghai_content and shenzhen_content:
        # 解析网页内容，获取数据
        shanghai_total_market_value, shanghai_circulating_market_value, shanghai_listed_companies, shanghai_average_pe_ratio = parse_page_content(shanghai_content, '上海交易所')
        shenzhen_total_market_value, shenzhen_circulating_market_value, shenzhen_listed_companies, shenzhen_average_pe_ratio = parse_page_content(shenzhen_content, '深圳交易所')

        # 保存不同交易所的数据到数据库
        save_data_to_db(today, '上海交易所', shanghai_total_market_value, shanghai_circulating_market_value, shanghai_listed_companies, shanghai_average_pe_ratio)
        save_data_to_db(today, '深圳交易所', shenzhen_total_market_value, shenzhen_circulating_market_value, shenzhen_listed_companies, shenzhen_average_pe_ratio)

if __name__ == '__main__':
    main()

# 关闭数据库连接
conn.close()